from flask import Flask, request
import torch
from exps.run11.net import ReversiNet

BOARD_SIZE = 15
model = ReversiNet().eval()
model.load_state_dict(torch.load('exps/run11/ema0001.pth', map_location='cpu'))

app = Flask(__name__)

@app.route("/step")
def step():
    board_str = request.args.get('board', '0'*225)
    board = torch.zeros((1, 2, BOARD_SIZE * BOARD_SIZE), dtype=bool)
    for i, b in enumerate(board_str):
        if b == '1':
            board[0, 0, i] = True
        if b == '2':
            board[0, 1, i] = True
    logits = model(board.reshape(1, 2, BOARD_SIZE, BOARD_SIZE).float())[0, 0]
    # logits.sigmoid().tolist()

    line = '<table>'
    for i, row in enumerate(logits.sigmoid().tolist()):
        line += '<tr>'
        for j, val in enumerate(row):
            newboard = [{'0': '0', '1': '2', '2': '1'}[b] for b in board_str]
            val = f'{val:.2f}'
            if newboard[i * BOARD_SIZE + j] == '1':
                val = 'o'
            elif newboard[i * BOARD_SIZE + j] == '2':
                val = 'x'
            else:
                newboard[i * BOARD_SIZE + j] = '2'
            newboard = ''.join(newboard)
            line += f'<td><a href="/step?board={newboard}">{val}</a></td>'
        line += '</tr>'
    line += '</table>'

    return line
